import threading
import openai
import os
import signal
import time
import queue
from tenacity import retry, stop_after_attempt, wait_random

#本程序仅在Windows11上测试，key是泰国号码注册，最好连接泰国节点
openai.api_key = "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"
q = queue.Queue()   #问题队列
identify = [0]      #记录线程号
count = 0           #记录用户数量

@retry(stop=stop_after_attempt(8), wait=wait_random(min=1, max=5)) #反复提交问题，这里指定提交6次,每次等待1-3秒。可根据自己情况修改。
def chat(prompt):
    try:
        response = openai.ChatCompletion.create(
          model="gpt-3.5-turbo",
          messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "Who won the world series in 2020?"},
                {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
                {"role": "user", "content": prompt}
            ]
        )
        answer = response['choices'][0]['message']['content'].strip()
        if len(answer) != 0:  #避免返回空白
            return answer
        else:
            raise Exception
    except Exception as e:
        print(e)

def ans(sequence):
    while True:
        if not q.empty():
            if q.queue[0][0] == threading.get_ident():
                text = ""  # 设置一个字符串变量
                turns = []  # 设置一个列表变量，turn指对话时的话轮
                question = q.get()[1]
                if len(question.strip()) == 0:  # 如果输入为空，提醒输入问题
                    print("please input your question")
                elif question.lower() == "quit":  # 如果输入为"quit"，本线程用户问答完成，线程结束
                    print("第" + str(sequence) + "位用户， See You Next Time ！")
                    break
                else:
                    prompt = text + "\n" + question
                    if len(prompt) <= 2000:  # 避免撑爆。ChatGpt API最大处理1500个词左右: prompt+completion不能超过2048个token，约1500个自然词。
                        result = chat(prompt)
                    else:
                        result = chat(prompt[-2000:])  # 因为len(prompt)算的是字符数,2000这个字符数可以自己调整，估计不超过5000一般都可以。
                    turns += [question] + [result]  # 只有这样迭代才能连续提问理解上下文
                    print("\n第" + str(sequence) + "位用户问题：" + question)
                    print("第" + str(sequence) + "位用户的答案", result)
                    if len(turns) <= 6:  # 指定一定的话轮语境以保证对话的连续性，这里指定为6次。你可以根据实际情况修改。
                        text = " ".join('%s' % id for id in (turns))
                    else:
                        text = " ".join(turns[-6:])



def worker(num):
    #线程的工作函数
    print("本线程序号为：" + str(threading.get_ident()))
    identify.append(threading.get_ident())
    ans(num)

def signal_handler(sig, frame):
    global count
    count += 1
    print('新增第' + str(count) + "用户")
    t = threading.Thread(target=worker, args=(count,))
    t.start()


# 注册信号处理函数
signal.signal(signal.SIGINT, signal_handler)

pid = os.getpid()
print("本进程pid为：" + str(pid))

while True:
    #在终端运行本程序，按Ctrl+c新增用户
    print("可新增用户")
    time.sleep(3)

    if q.empty() and len(identify) > 1:
        print("不可新增用户！！！！！")
        print("现有" + str(len(identify)-1) + "位用户")
        i = input("输入提问用户序号：")
        if i.isdigit():
            print(i)
            print(len(identify)-1)
            if int(i) <= (len(identify)-1):
                j = identify[int(i)]
                s = input("输入第" + str(i) + "位用户的问题：")
                my_tuple = (j, s)
                q.put(my_tuple)
            else:
                print("输入序号超出范围")
        else:
            print("输入不为数字")